# preprocess.py
import os
from collections import Counter
import torch
from utils.vocab import Vocab

def build_vocab(file_path, vocab_path, min_freq=1):  # 设置最小频率阈值为 1
    with open(file_path, 'r', encoding='utf-8') as f:
        words = [word for line in f for word in line.strip().split()]
    counter = Counter(words)
    vocab = {'<pad>': 0, '<sos>': 1, '<eos>': 2, '<unk>': 3}
    itos = ['<pad>', '<sos>', '<eos>', '<unk>']
    for word, freq in counter.items():
        if freq >= min_freq:
            vocab[word] = len(vocab)
            itos.append(word)
    torch.save({'stoi': vocab, 'itos': itos}, vocab_path)

if __name__ == "__main__":
    data_dir = './data'
    build_vocab(os.path.join(data_dir, 'TM-training-set/chinese.txt'), os.path.join(data_dir, 'src_vocab.pth'), min_freq=1)
    build_vocab(os.path.join(data_dir, 'TM-training-set/english.txt'), os.path.join(data_dir, 'tgt_vocab.pth'), min_freq=1)